package edu.nepu.flink.api.combinStream;

import edu.nepu.flink.api.bean.WaterSensor;
import org.apache.flink.streaming.api.datastream.ConnectedStreams;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.co.CoMapFunction;
import org.apache.flink.streaming.api.functions.co.CoProcessFunction;
import org.apache.flink.util.Collector;
import org.apache.kafka.common.protocol.types.Field;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * @Date 2024/2/28 22:37
 * @Created by chenshuaijun
 */
public class ConnectDemo {

    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(2);
        DataStreamSource<WaterSensor> streamSource1 = env.fromElements(
                new WaterSensor("s1",1L,1),
                new WaterSensor("s1",2L,2),
                new WaterSensor("s2",3L,3)
        );
        DataStreamSource<WaterSensor> streamSource2 = env.fromElements(
                new WaterSensor("s1",2L,2),
                new WaterSensor("s2",4L,4)
        );

        ConnectedStreams<WaterSensor, WaterSensor> connectStream = streamSource1.connect(streamSource2);

        SingleOutputStreamOperator<String> processStream = connectStream.process(new CoProcessFunction<WaterSensor, WaterSensor, String>() {

            Map<String, List<WaterSensor>> s1Map = new HashMap<>();
            Map<String, List<WaterSensor>> s2Map = new HashMap<>();

            @Override
            public void processElement1(WaterSensor value, CoProcessFunction<WaterSensor, WaterSensor, String>.Context ctx, Collector<String> out) throws Exception {
                String id = value.getId();
                List<WaterSensor> orDefault = s1Map.getOrDefault(id, new ArrayList<>());
                orDefault.add(value);
                s1Map.put(id, orDefault);
                if (s2Map.containsKey(id)) {
                    List<WaterSensor> waterSensors = s2Map.get(id);
                    for (WaterSensor waterSensor : waterSensors) {
                        out.collect(waterSensor.toString() + "->" + value.toString());
                    }
                }
            }

            @Override
            public void processElement2(WaterSensor value, CoProcessFunction<WaterSensor, WaterSensor, String>.Context ctx, Collector<String> out) throws Exception {
                String id = value.getId();
                List<WaterSensor> orDefault = s2Map.getOrDefault(id, new ArrayList<>());
                orDefault.add(value);
                s2Map.put(id, orDefault);
                if (s1Map.containsKey(id)) {
                    List<WaterSensor> waterSensors = s1Map.get(id);
                    for (WaterSensor waterSensor : waterSensors) {
                        out.collect(waterSensor.toString() + "->" + value.toString());
                    }
                }
            }
        });

        processStream.print();

        env.execute();

    }
}
